package tree

import (
	"encoding/json"
	"fmt"
)

type (
	NodeBase interface {
		GetChildren() []NodeBase
		SetChildren([]NodeBase)
	}
	NodeWithChildren interface {
		GetID() int64
		GetParentID() int64
		GetChildren() []NodeWithChildren
		SetChildren([]NodeWithChildren)
	}
	NodeWithCode interface {
		GetCode() string
		GetParentCode() string
		GetChildren() []NodeWithCode
		SetChildren([]NodeWithCode)
	}
)

// Implement the buildTreeByID function
func buildTreeByID(nodes []NodeWithChildren) []NodeWithChildren {
	nodeMap := make(map[int64]NodeWithChildren)
	for _, node := range nodes {
		nodeMap[node.GetID()] = node
	}
	rootNodes := make([]NodeWithChildren, 0)
	for _, node := range nodes {
		parentID := node.GetParentID()
		if parentID == 0 {
			rootNodes = append(rootNodes, node)
		} else {
			parent, ok := nodeMap[parentID]
			if ok {
				// 获取父节点的子节点列表
				children := parent.GetChildren()

				// 确保子节点列表中没有重复添加当前节点
				found := false
				for _, child := range children {
					if child.GetID() == node.GetID() {
						found = true
						break
					}
				}

				// 如果没有找到，则添加当前节点到父节点的children中
				if !found {
					children = append(children, node)
					parent.SetChildren(children)
				}
			} else {
				// 如果找不到父节点，则把当前节点作为根节点
				rootNodes = append(rootNodes, node)
			}
		}
	}
	return rootNodes
}

func buildTreeByCode(nodes []NodeWithCode) []NodeWithCode {
	nodeMap := make(map[string]NodeWithCode)
	for _, node := range nodes {
		nodeMap[node.GetCode()] = node
	}
	rootNodes := make([]NodeWithCode, 0)
	for _, node := range nodes {
		parentCode := node.GetParentCode()
		if parentCode == "" {
			rootNodes = append(rootNodes, node)
		} else {
			parent, ok := nodeMap[parentCode]
			if ok {
				children := parent.GetChildren()
				children = append(children, node.(NodeWithCode))
				parent.SetChildren(children)
			} else {
				rootNodes = append(rootNodes, node)
			}

		}
	}
	return rootNodes
}

func MakeTreeByID(records []NodeWithChildren) []NodeWithChildren {
	return buildTreeByID(records)
}

func MakeTreeByCode(records []NodeWithCode) []NodeWithCode {
	return buildTreeByCode(records)
}

func PrintTree(trees interface{}) {
	b, _ := json.MarshalIndent(trees, "", "    ")
	fmt.Println(string(b))
}
